import os
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
import imageio
import tqdm

def get_mean_std(pathDir: list):
    # 图像尺寸
    img = imageio.imread(os.path.join(filepath, pathDir[0]))
    num = len(pathDir) * img.shape[0] * img.shape[1]

    # 计算三通道的均值
    R_channel = 0
    G_channel = 0
    B_channel = 0
    for idx in tqdm.tgrange(len(pathDir)):
        filename = pathDir[idx]
        img = imageio.imread(os.path.join(filepath, filename))
        R_channel = R_channel + np.sum(img[:, :, 0])
        G_channel = G_channel + np.sum(img[:, :, 1])
        B_channel = B_channel + np.sum(img[:, :, 2])

    R_mean = R_channel / num
    G_mean = G_channel / num
    B_mean = B_channel / num

    # 计算三通道的标准差
    R_channel = 0
    G_channel = 0
    B_channel = 0
    for idx in tqdm.tgrange(len(pathDir)):
        filename = pathDir[idx]
        img = imageio.imread(os.path.join(filepath, filename))
        R_channel = R_channel + np.sum((img[:, :, 0] - R_mean) ** 2)
        G_channel = G_channel + np.sum((img[:, :, 1] - G_mean) ** 2)
        B_channel = B_channel + np.sum((img[:, :, 2] - B_mean) ** 2)

    R_std = np.sqrt(R_channel / num)
    G_std = np.sqrt(G_channel / num)
    B_std = np.sqrt(B_channel / num)

    print("R_mean is %f, G_mean is %f, B_mean is %f" % (R_mean, G_mean, B_mean))
    print("R_var is %f, G_var is %f, B_var is %f" % (R_std, G_std, B_std))

    print("R_mean is %f, G_mean is %f, B_mean is %f" % (R_mean/255.0, G_mean/255.0, B_mean/255.0))
    print("R_var is %f, G_var is %f, B_var is %f" % (R_std/255.0, G_std/255.0, B_std/255.0))
    
    return (R_mean, G_mean, B_mean), (R_std, G_std, B_std)

if __name__ == "__main__":
    filepath = '/root/ws/data/meshb_qgy/processed_dir/train/images/'  # 数据集目录
    pathDir = os.listdir(filepath)
    get_mean_std(pathDir)
